##计算ATR数值（移植自主程/波动率.py，移除对主程路径的依赖）
from typing import Optional, Dict
from datetime import datetime
import time
import pandas as pd
import numpy as np

try:
    import talib as ta
except Exception:
    ta = None

from xtquant import xtdata


class ATRCalculator:
    """
    ATR 波动率计算器（基于当日 tick 重采样为 1 分钟K）
    - TR = max(high-low, |high-prev_close|, |low-prev_close|)
    - ATR 计算方式：优先 TA-Lib（Wilder），否则回退 EMA/SMA
    - 内置简单缓存，避免每个tick都重复重算
    """

    def __init__(self, data_provider=xtdata, lookback: int = 14, method: str = "talib",
                 cache_ttl_seconds: int = 10, verbose: bool = False):
        self._xt = data_provider
        self.lookback = int(lookback)
        self.method = method.lower().strip()  # "talib" | "ema" | "sma"
        self.cache_ttl_seconds = int(cache_ttl_seconds)
        self.verbose = bool(verbose)
        self._cache: Dict[str, Dict[str, float]] = {}

    def _log(self, msg: str):
        if self.verbose:
            print(f"【ATR】{msg}")

    def _now(self) -> float:
        return time.time()

    def _download_today_ticks(self, stock_code: str) -> Optional[pd.DataFrame]:
        try:
            today = datetime.now().strftime("%Y%m%d")
            self._xt.download_history_data(stock_code=stock_code, period="tick", start_time=today, end_time=today)
            time.sleep(1.0)
            data = self._xt.get_local_data(stock_list=[stock_code], period='tick', start_time=today, end_time=today)
            if data and stock_code in data:
                return data[stock_code]
            return None
        except Exception:
            return None

    def _ensure_time_and_price(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
        if df is None or df.empty:
            return None
        price_col = next((col for col in ['price', 'lastPrice', 'close', 'last', 'current'] if col in df.columns), None)
        if not price_col:
            return None
        time_col = next((col for col in df.columns if 'time' in col.lower()), None)
        if not time_col:
            return None
        try:
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['time'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                df['time'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
            df = df[df['time'].notna()].sort_values('time')
        except Exception:
            return None
        if df.empty:
            return None
        df = df[df[price_col] > 0].copy()
        lunch_start = pd.Timestamp('11:30:00').time()
        lunch_end = pd.Timestamp('12:59:59').time()
        df['time_only'] = df['time'].dt.time
        df = df[~((df['time_only'] >= lunch_start) & (df['time_only'] <= lunch_end))]
        if df.empty:
            return None
        df.rename(columns={price_col: 'price'}, inplace=True)
        return df

    def _resample_to_minutes(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
        try:
            df = df.set_index('time')
            ohlc = df['price'].resample('1min').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
            }).dropna()
            ohlc = ohlc.iloc[-self.lookback:]
            return ohlc.reset_index()
        except Exception:
            return None

    def _compute_tr(self, df_minutes: pd.DataFrame) -> Optional[pd.Series]:
        if df_minutes is None or df_minutes.empty:
            return None
        high = df_minutes['high'].astype(float)
        low = df_minutes['low'].astype(float)
        close = df_minutes['close'].astype(float)
        prev_close = close.shift(1)
        tr = pd.concat([
            (high - low),
            (high - prev_close).abs(),
            (low - prev_close).abs()
        ], axis=1).max(axis=1)
        return tr

    def _ema(self, series: pd.Series, period: int) -> pd.Series:
        if period <= 0 or series is None or series.empty:
            return pd.Series(dtype=float)
        alpha = 2 / (period + 1)
        ema_vals = []
        last = None
        for val in series:
            if last is None:
                last = val
            else:
                last = alpha * val + (1 - alpha) * last
            ema_vals.append(last)
        return pd.Series(ema_vals, index=series.index)

    def _sma(self, series: pd.Series, period: int) -> pd.Series:
        try:
            return series.rolling(window=period, min_periods=period).mean()
        except Exception:
            return pd.Series(dtype=float)

    def _compute_atr(self, df_minutes: pd.DataFrame, lookback: int) -> Optional[float]:
        tr = self._compute_tr(df_minutes)
        if tr is None or tr.empty:
            return None
        current_price = df_minutes['close'].iloc[-1] if not df_minutes.empty else None
        tr = tr.iloc[-lookback:]
        if len(tr) < lookback:
            return None
        if current_price is not None and any(tr > current_price * 0.1):
            return None
        method_lower = self.method.lower()
        if method_lower == 'talib' and ta is not None:
            try:
                atr_values = ta.ATR(
                    df_minutes['high'].astype(float).values,
                    df_minutes['low'].astype(float).values,
                    df_minutes['close'].astype(float).values,
                    timeperiod=lookback
                )
            except Exception:
                return None
            if atr_values is None or len(atr_values) == 0 or np.isnan(atr_values[-1]):
                return None
            return float(atr_values[-1])
        if method_lower == 'ema':
            atr_series = self._ema(tr, lookback)
        else:
            atr_series = self._sma(tr, lookback)
        return float(atr_series.iloc[-1]) if not atr_series.empty else None

    def calculate_atr(self, stock_code: str, current_price: Optional[float] = None, lookback: Optional[int] = None) -> Optional[float]:
        lookback = int(lookback or self.lookback)
        cached = self._cache.get(stock_code)
        if cached and (self._now() - cached.get('ts', 0)) < self.cache_ttl_seconds:
            return cached.get('atr')
        ticks_df = self._download_today_ticks(stock_code)
        ticks_df = self._ensure_time_and_price(ticks_df)
        if ticks_df is None:
            self._log(f"{stock_code} 无有效tick数据，ATR返回None")
            return None
        minutes_df = self._resample_to_minutes(ticks_df)
        if minutes_df is None or minutes_df.empty:
            self._log(f"{stock_code} 分钟数据为空，ATR返回None")
            return None
        atr_value = self._compute_atr(minutes_df, lookback)
        if atr_value is None:
            self._log(f"{stock_code} ATR计算失败")
            return None
        self._cache[stock_code] = {"ts": self._now(), "atr": atr_value}
        return atr_value


